import socket
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding

def decrypt(ciphertext, key, iv):
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    decryptor = cipher.decryptor()
    padded_data = decryptor.update(ciphertext) + decryptor.finalize()
    unpadder = padding.PKCS7(128).unpadder()
    return unpadder.update(padded_data) + unpadder.finalize()

# 配置参数
key = b'this_is_a_32_byte_key_for_aes256!!'  # 必须与客户端一致
output_path = 'output.txt'                   # 保存解密后的文件

# 建立服务器
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_socket.bind(('localhost', 9999))
server_socket.listen(5)
print("等待连接...")

client_socket, addr = server_socket.accept()
print(f"来自 {addr} 的连接已建立")

# 接收IV（16字节）
iv = client_socket.recv(16)

# 接收数据长度（4字节）
data_len = int.from_bytes(client_socket.recv(4), 'big')

# 接收加密数据
ciphertext = b''
while len(ciphertext) < data_len:
    chunk = client_socket.recv(min(data_len - len(ciphertext), 4096))
    if not chunk:
        break
    ciphertext += chunk

# 解密并保存文件
try:
    plaintext = decrypt(ciphertext, key, iv)
    with open(output_path, 'wb') as f:
        f.write(plaintext)
    client_socket.send(b"File received and decrypted successfully!")
except Exception as e:
    client_socket.send(f"Error: {str(e)}".encode())

client_socket.close()
server_socket.close()